;;---------------------------------------------------------------------------
;; prepare data for 
;; 
;; . zonal mean 
;; . surface 
;; . on pressure levels 
;;
;; kai zhang, max planck institute for meteorology 
;;  
;;---------------------------------------------------------------------------

load "plot_lib.ncl" 


begin 

  data_class= "burden" 
  data_class= "mass" 
  data_class= "num" 
  data_class= "aod" 

  data_type = "zonal" 
  data_type = "surf" 

;; experiment name 

  casa = "E0001B" 
  casb = "HAM2R46H" 

  ddir = "./" 

  load "load.ncl" 

  if(data_class.eq."burden") then 
     data_type = "surf" 
  end if 

  if(data_class.eq."aod") then
     data_type = "surf"
  end if

  if(data_type.eq."zonal") then 
     ftyp = "P" 
  end if 

  if(data_type.eq."surf") then 
     ftyp = "S" 
  end if 

  if(data_class.eq."burden") then
     ftyp = "BD"
  end if

  if(data_class.eq."aod") then
     ftyp = ""
  end if


;; experiment year 

  year = "2000" 

;; data directory 

  path = "./"
  path = "/scratch/work/bm0103/m222044/exp/diag/"
  pata = path+casa+"/" 
  patb = path+casb+"/" 

;; filename 

  if(data_class.eq."mass".or.data_class.eq."burden") then 

     fnaa = pata+casa+".MSU"+ftyp+"."+year+".echam.nc"
     fnab = pata+casa+".MBC"+ftyp+"."+year+".echam.nc"
     fnac = pata+casa+".MPM"+ftyp+"."+year+".echam.nc"
     fnad = pata+casa+".MSO"+ftyp+"."+year+".echam.nc"
     fnae = pata+casa+".MDU"+ftyp+"."+year+".echam.nc"
     fnaf = pata+casa+".MSS"+ftyp+"."+year+".echam.nc"
     fnag = pata+casa+".DMS"+ftyp+"."+year+".echam.nc"
     fnah = pata+casa+".SO2"+ftyp+"."+year+".echam.nc"
     fnai = pata+casa+".SO4"+ftyp+"."+year+".echam.nc"

     fnba = patb+casb+".MSU"+ftyp+"."+year+".echam.nc"
     fnbb = patb+casb+".MBC"+ftyp+"."+year+".echam.nc"
     fnbc = patb+casb+".MPM"+ftyp+"."+year+".echam.nc"
     fnbd = patb+casb+".MSO"+ftyp+"."+year+".echam.nc"
     fnbe = patb+casb+".MDU"+ftyp+"."+year+".echam.nc"
     fnbf = patb+casb+".MSS"+ftyp+"."+year+".echam.nc"
     fnbg = patb+casb+".DMS"+ftyp+"."+year+".echam.nc"
     fnbh = patb+casb+".SO2"+ftyp+"."+year+".echam.nc"
     fnbi = patb+casb+".SO4"+ftyp+"."+year+".echam.nc"

  end if 

  if(data_class.eq."mass") then
     fnai = pata+casa+".SO4_gas"+ftyp+"."+year+".echam.nc"
     fnbi = patb+casb+".SO4_gas"+ftyp+"."+year+".echam.nc"
  end if

  if(data_class.eq."num") then

     fnaa = pata+casa+".NUM_NS"+ftyp+"."+year+".echam.nc"
     fnab = pata+casa+".NUM_KS"+ftyp+"."+year+".echam.nc"
     fnac = pata+casa+".NUM_AS"+ftyp+"."+year+".echam.nc"
     fnad = pata+casa+".NUM_CS"+ftyp+"."+year+".echam.nc"
     fnae = pata+casa+".NUM_KI"+ftyp+"."+year+".echam.nc"
     fnaf = pata+casa+".NUM_AI"+ftyp+"."+year+".echam.nc"
     fnag = pata+casa+".NUM_CI"+ftyp+"."+year+".echam.nc"
     fnah = pata+casa+".NUM_NS"+ftyp+"."+year+".echam.nc"
     fnai = "nothing" ;;pata+casa+".NUM_NS"+ftyp+"."+year+".echam.nc"

     fnba = patb+casb+".NUM_NS"+ftyp+"."+year+".echam.nc"
     fnbb = patb+casb+".NUM_KS"+ftyp+"."+year+".echam.nc"
     fnbc = patb+casb+".NUM_AS"+ftyp+"."+year+".echam.nc"
     fnbd = patb+casb+".NUM_CS"+ftyp+"."+year+".echam.nc"
     fnbe = patb+casb+".NUM_KI"+ftyp+"."+year+".echam.nc"
     fnbf = patb+casb+".NUM_AI"+ftyp+"."+year+".echam.nc"
     fnbg = patb+casb+".NUM_CI"+ftyp+"."+year+".echam.nc"
     fnbh = patb+casb+".NUM_NS"+ftyp+"."+year+".echam.nc"
     fnbi = "nothing" ;;patb+casb+".NUM_NS"+ftyp+"."+year+".echam.nc"

  end if

  if(data_class.eq."aod") then

     fnaa = pata+casa+".radmm"+ftyp+"."+year+".echam.nc"
     fnab = pata+casa+".radmm"+ftyp+"."+year+".echam.nc"
     fnac = pata+casa+".radmm"+ftyp+"."+year+".echam.nc"
     fnad = "nothing" 
     fnae = "nothing" 
     fnaf = "nothing" 
     fnag = "nothing" 
     fnah = "nothing" 
     fnai = "nothing" 

     fnba = patb+casb+".radmm"+ftyp+"."+year+".echam.nc"
     fnbb = patb+casb+".radmm"+ftyp+"."+year+".echam.nc"
     fnbc = patb+casb+".radmm"+ftyp+"."+year+".echam.nc"
     fnbd = "nothing" 
     fnbe = "nothing" 
     fnbf = "nothing" 
     fnbg = "nothing" 
     fnbh = "nothing" 
     fnbi = "nothing" 

  end if

  print("% open files: ")
  print("% ")
  print("% ")
  print("% "+fnaa)
  print("% "+fnab)
  print("% "+fnac)
  print("% "+fnad)
  print("% "+fnae)
  print("% "+fnaf)
  print("% "+fnag)
  print("% "+fnah)
  print("% "+fnai)
  print("% ")
  print("% ")
  print("% "+fnba)
  print("% "+fnbb)
  print("% "+fnbc)
  print("% "+fnbd)
  print("% "+fnbe)
  print("% "+fnbf)
  print("% "+fnbg)
  print("% "+fnbh)
  print("% "+fnbi)
  print("% ")
  print("% ")

  flaa = sysopenfile(fnaa) ;; see plot_lib 
  flab = sysopenfile(fnab) 
  flac = sysopenfile(fnac)
  flad = sysopenfile(fnad)
  flae = sysopenfile(fnae)
  flaf = sysopenfile(fnaf)
  flag = sysopenfile(fnag)
  flah = sysopenfile(fnah)
  flai = sysopenfile(fnai)

  flba = sysopenfile(fnba)
  flbb = sysopenfile(fnbb)
  flbc = sysopenfile(fnbc)
  flbd = sysopenfile(fnbd)
  flbe = sysopenfile(fnbe)
  flbf = sysopenfile(fnbf)
  flbg = sysopenfile(fnbg)
  flbh = sysopenfile(fnbh)
  flbi = sysopenfile(fnbi)

;; the following part reads the data 
;;
;; time = 12
;; lev  = 14 
;; lat  = 96 
;; lon  = 192 

  print("% reading data from 1st experiment") 
  print("% ") 
  print("% ") 
;; data from 1st experiment  

  if(data_class.eq."burden") then 
     msua = flaa->MSUBD
     mbca = flab->MBCBD
     mpma = flac->MPMBD
     msoa = flac->MPMBD
     msoa = 0.
     mdua = flae->MDUBD
     mssa = flaf->MSSBD
     dmsa = flag->D_BURDEN_DMS
     so2a = flah->D_BURDEN_SO2
     so4a = flai->D_BURDEN_SO4_gas
     vbas = msua 
  end if
  if(data_class.eq."mass") then 
     msua = flaa->var ;;MSU  
     mbca = flab->var ;;MBC  
     mpma = flac->var ;;MPM  
     msoa = flac->var ;;MPM  
     msoa = 0. 
     mdua = flae->var ;;MDU  
     mssa = flaf->var ;;MSS  
     dmsa = flag->var ;;DMS
     so2a = flah->var ;;SO2  
     so4a = flai->var ;;SO4  
     vbas = msua 
  end if 
  if(data_class.eq."num") then
     nusa = flaa->var ;;
     nksa = flab->var ;;
     nasa = flac->var ;;
     ncsa = flad->var ;;
     nkia = flae->var ;;
     naia = flaf->var ;;
     ncia = flag->var ;;
     vbas = nusa 
  end if
  if(data_class.eq."aod") then
     taoa = flaa->TAU_2D ;;
     aaoa = flab->ABS_2D ;;
     anga = flac->ANG    ;;
     vbas = taoa
  end if

  print("% reading data from 2nd experiment") 
  print("% ") 
  print("% ") 
;; data from 2nd experiment  

  if(data_class.eq."burden") then 
     msub = flba->MSUBD
     mbcb = flbb->MBCBD
     mpmb = flbc->MPMBD
     msob = flbd->MSOBD
     mpmb = mpmb + msob
     mdub = flbe->MDUBD
     mssb = flbf->MSSBD
     dmsb = flbg->D_BURDEN_DMS
     so2b = flbh->D_BURDEN_SO2
     so4b = flbi->D_BURDEN_SO4_gas
     vbas = msub 
  end if 
  if(data_class.eq."mass") then 
     msub = flba->var ;;MSU  
     mbcb = flbb->var ;;MBC  
     mpmb = flbc->var ;;MPM  
     msob = flbd->var ;;MSO  
     mpmb = mpmb + msob 
     mdub = flbe->var ;;MDU  
     mssb = flbf->var ;;MSS  
     dmsb = flbg->var ;;DMS
     so2b = flbh->var ;;SO2  
     so4b = flbi->var ;;SO2  
     vbas = msub 
  end if 
  if(data_class.eq."num") then
     nusb = flba->var ;;
     nksb = flbb->var ;;
     nasb = flbc->var ;;
     ncsb = flbd->var ;;
     nkib = flbe->var ;;
     naib = flbf->var ;;
     ncib = flbg->var ;;
     vbas = nusb 
  end if
  if(data_class.eq."aod") then
     taob = flba->TAU_2D ;;
     aaob = flbb->ABS_2D ;;
     angb = flbc->ANG    ;;
     vbas = taob
  end if

  if(data_type.eq."zonal") then 
     nlon = dimsizes(vbas(0,0,0,:))
     nlat = dimsizes(vbas(0,0,:,0))
     nlev = dimsizes(vbas(0,:,0,0))
     ntim = dimsizes(vbas(:,0,0,0))
  end if 

  if(data_type.eq."surf") then
     nlon = dimsizes(vbas(0,0,:))
     nlat = dimsizes(vbas(0,:,0))
     ntim = dimsizes(vbas(:,0,0))
     nlev = -999 
  end if
 
  if(data_class.eq."mass" .or. data_class.eq."burden") then
     nsp = 9 
     imsu = 0 
     imbc = 1 
     impm = 2 
     imso = 3 
     imdu = 4 
     imss = 5 
     idms = 6 
     iso2 = 7 
     iso4 = 8 
     varname_list = (/"msu","mbc","mpm","mso","mdu","mss","dms","so2","so4"/) 
  end if 

  if(data_class.eq."num") then
     nsp = 7 
     inus = 0 
     inks = 1 
     inas = 2 
     incs = 3 
     inki = 4 
     inai = 5 
     inci = 6 
     varname_list = (/"nus","nks","nas","ncs","nki","nai","nci"/) 
  end if 

  if(data_class.eq."aod") then
     nsp = 3
     itao = 0
     iaao = 1
     iang = 2
     varname_list = (/"TAU_2D","ABS_2D","ANG"/)
  end if

;;----------------------------------------------------------------------
;; define the variable 
;;----------------------------------------------------------------------
  
  if(data_type.eq."zonal") then 
     vpa = new((/nsp,ntim,nlev,nlat,nlon/),"float") 
     vpa!0 = "vid"
     vpa!1 = "time" 
     vpa!2 = "lev" 
     vpa!3 = "lat" 
     vpa!4 = "lon" 
  end if 

  if(data_type.eq."surf") then
     vpa = new((/nsp,ntim,nlat,nlon/),"float")
     vpa!0 = "vid"
     vpa!1 = "time"
     vpa!2 = "lat"
     vpa!3 = "lon"
  end if

  vpb = vpa
  vpc = vpa

;;----------------------------------------------------------------------
;; build up the variable 
;;----------------------------------------------------------------------
 
  if(data_type.eq."zonal" .and. data_class.eq."mass") then 
     vpa(imsu,:,:,:,:) = (/ msua /)  
     vpa(imbc,:,:,:,:) = (/ mbca /)  
     vpa(impm,:,:,:,:) = (/ mpma /)  
     vpa(imso,:,:,:,:) = (/ msoa /)  
     vpa(imdu,:,:,:,:) = (/ mdua /)  
     vpa(imss,:,:,:,:) = (/ mssa /)  
     vpa(idms,:,:,:,:) = (/ dmsa /)  
     vpa(iso2,:,:,:,:) = (/ so2a /)  
     vpa(iso4,:,:,:,:) = (/ so4a /)  
 
     vpb(imsu,:,:,:,:) = (/ msub /)  
     vpb(imbc,:,:,:,:) = (/ mbcb /)  
     vpb(impm,:,:,:,:) = (/ mpmb /)  
     vpb(imso,:,:,:,:) = (/ msob /)  
     vpb(imdu,:,:,:,:) = (/ mdub /)  
     vpb(imss,:,:,:,:) = (/ mssb /)  
     vpb(idms,:,:,:,:) = (/ dmsb /)  
     vpb(iso2,:,:,:,:) = (/ so2b /)  
     vpb(iso4,:,:,:,:) = (/ so4b /)  
  end if 

  if(data_type.eq."surf" .and. (data_class.eq."mass" .or. data_class.eq."burden") ) then
     vpa(imsu,:,:,:) = (/ msua /)
     vpa(imbc,:,:,:) = (/ mbca /)
     vpa(impm,:,:,:) = (/ mpma /)
     vpa(imso,:,:,:) = (/ msoa /)
     vpa(imdu,:,:,:) = (/ mdua /)
     vpa(imss,:,:,:) = (/ mssa /)
     vpa(idms,:,:,:) = (/ dmsa /)
     vpa(iso2,:,:,:) = (/ so2a /)
     vpa(iso4,:,:,:) = (/ so4a /)

     vpb(imsu,:,:,:) = (/ msub /)
     vpb(imbc,:,:,:) = (/ mbcb /)
     vpb(impm,:,:,:) = (/ mpmb /)
     vpb(imso,:,:,:) = (/ msob /)
     vpb(imdu,:,:,:) = (/ mdub /)
     vpb(imss,:,:,:) = (/ mssb /)
     vpb(idms,:,:,:) = (/ dmsb /)
     vpb(iso2,:,:,:) = (/ so2b /)
     vpb(iso4,:,:,:) = (/ so4b /)
  end if

  if(data_type.eq."zonal" .and. data_class.eq."num") then
     vpa(inus,:,:,:,:) = (/ nusa /)
     vpa(inks,:,:,:,:) = (/ nksa /)
     vpa(inas,:,:,:,:) = (/ nasa /)
     vpa(incs,:,:,:,:) = (/ ncsa /)
     vpa(inki,:,:,:,:) = (/ nkia /)
     vpa(inai,:,:,:,:) = (/ naia /)
     vpa(inci,:,:,:,:) = (/ ncia /)

     vpb(inus,:,:,:,:) = (/ nusb /)
     vpb(inks,:,:,:,:) = (/ nksb /)
     vpb(inas,:,:,:,:) = (/ nasb /)
     vpb(incs,:,:,:,:) = (/ ncsb /)
     vpb(inki,:,:,:,:) = (/ nkib /)
     vpb(inai,:,:,:,:) = (/ naib /)
     vpb(inci,:,:,:,:) = (/ ncib /)
  end if

  if(data_type.eq."surf" .and. data_class.eq."num") then
     vpa(inus,:,:,:) = (/ nusa /)
     vpa(inks,:,:,:) = (/ nksa /)
     vpa(inas,:,:,:) = (/ nasa /)
     vpa(incs,:,:,:) = (/ ncsa /)
     vpa(inki,:,:,:) = (/ nkia /)
     vpa(inai,:,:,:) = (/ naia /)
     vpa(inci,:,:,:) = (/ ncia /)

     vpb(inus,:,:,:) = (/ nusb /)
     vpb(inks,:,:,:) = (/ nksb /)
     vpb(inas,:,:,:) = (/ nasb /)
     vpb(incs,:,:,:) = (/ ncsb /)
     vpb(inki,:,:,:) = (/ nkib /)
     vpb(inai,:,:,:) = (/ naib /)
     vpb(inci,:,:,:) = (/ ncib /)
  end if

  if(data_class.eq."aod") then
     vpa(itao,:,:,:) = (/ taoa /)
     vpa(iaao,:,:,:) = (/ aaoa /)
     vpa(iang,:,:,:) = (/ anga /)

     vpb(itao,:,:,:) = (/ taob /)
     vpb(iaao,:,:,:) = (/ aaob /)
     vpb(iang,:,:,:) = (/ angb /)
  end if
 
  ;; difference 

  vpc = (/ vpb - vpa /)  

;;----------------------------------------------------------------------
;; zonal 
;;----------------------------------------------------------------------

  if(data_type.eq."zonal") then 

     ;; zonal mean 

     zvpa = vpa(:,:,:,:,0) 
     zvpa!0 = "vid"
     zvpa!1 = "time"
     zvpa!2 = "lev"
     zvpa!3 = "lat"
     zvpb = zvpa 
     zvpc = zvpa 

     zvpa = (/ xdimavg(vpa,"lon") /)  
     zvpb = (/ xdimavg(vpb,"lon") /)  
     zvpc = (/ xdimavg(vpc,"lon") /)  

     ;; annual mean

     va = vpa(:,0,:,:,0) 
     va!0       = "vid" 
     va&vid     = ispan(1,nsp,1) 
     va!1       = "lev" 
     va&lev     = vbas&lev
     va!2       = "lat" 
     va&lat     = vbas&lat
     
     vb = va
     vc = va

     va = (/ xdimavg(zvpa,"time") /)  
     vb = (/ xdimavg(zvpb,"time") /)  
     vc = (/ xdimavg(zvpc,"time") /)  

  end if 

;;----------------------------------------------------------------------
;; surf  
;;----------------------------------------------------------------------

  if(data_type.eq."surf") then

     va         = vpa(:,0,:,:)
     va!0       = "vid"
     va&vid     = ispan(1,nsp,1)
     va!1       = "lat"
     va&lat     = vbas&lat
     va!2       = "lon"
     va&lon     = vbas&lon
     vb         = va
     vc         = va

     va = (/ xdimavg(vpa,"time") /) ;; annual mean 
     vb = (/ xdimavg(vpb,"time") /)
     vc = (/ xdimavg(vpc,"time") /)

  end if


  varname_list_char = stringtochar(varname_list) 

  fno = ddir+"input_"+casa+"_"+casb+"_"+data_class+"_"+data_type+".nc" 

  print("% output file:  "+fno)
  print("% ")
  print("% ")

  varlist = va(:,0,0) 

  if(data_class.eq."mass".or.data_class.eq."burden") then
     varlist@var_09 = varname_list(8) 
     varlist@var_08 = varname_list(7) 
     varlist@var_07 = varname_list(6)
     varlist@var_06 = varname_list(5)
     varlist@var_05 = varname_list(4)
     varlist@var_04 = varname_list(3)
     varlist@var_03 = varname_list(2)
     varlist@var_02 = varname_list(1)
     varlist@var_01 = varname_list(0)
  end if

  if(data_class.eq."aod") then 
     varlist@var_03 = varname_list(2)
     varlist@var_02 = varname_list(1)
     varlist@var_01 = varname_list(0)
  end if 

  if(data_class.eq."num") then 
     varlist@var_07 = varname_list(6) 
     varlist@var_06 = varname_list(5) 
     varlist@var_05 = varname_list(4) 
     varlist@var_04 = varname_list(3) 
     varlist@var_03 = varname_list(2) 
     varlist@var_02 = varname_list(1) 
     varlist@var_01 = varname_list(0) 
  end if 

  if(data_type.eq."zonal") then  
     delete(varlist@lev) 
     delete(varlist@lat) 
     delete(varlist@_FillValue) 
  end if 
  if(data_type.eq."surf") then  
     delete(varlist@lon) 
     delete(varlist@lat) 
     delete(varlist@_FillValue) 
  end if 


  system("rm "+fno) 

  flo = addfile(fno,"c") 

  flo->ctrl=va
  flo->test=vb
  flo->diff=vc
  flo->varlist=varlist 

  print("% ------ end of program ------ ")
  print("% ")
  print("% ")
  print("% ")
  print("% ")

end






